import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from collections import Counter
from pinecone import Pinecone, ServerlessSpec

# 1. 初始化 Pinecone
def initialize_pinecone(api_key):
    """初始化Pinecone客户端"""
    try:
        pc = Pinecone(api_key=api_key)
        print("Pinecone初始化成功")
        return pc
    except Exception as e:
        print(f"Pinecone初始化失败: {e}")
        return None

# 2. 创建并连接到索引
def create_and_connect_index(pc, index_name="mnist-index"):
    """创建Pinecone索引并连接"""
    
    # 检查、清除之前创建的索引
    existing_indexes = pc.list_indexes()
    
    if any(index_info.name == index_name for index_info in existing_indexes):
        print(f"索引 '{index_name}' 已存在，正在删除...")
        pc.delete_index(index_name)
        print(f"索引 '{index_name}' 已成功删除...")
    else:
        print(f"索引 '{index_name}' 不存在，将创建新索引...")
    
    # 创建新索引
    print(f"正在创建新索引 '{index_name}' ...")
    pc.create_index(
        name=index_name,
        dimension=64,  # MNIST每个图像展平后是64维向量
        metric="euclidean",  # 使用欧氏距离
        spec=ServerlessSpec(
            cloud="aws",
            region="us-east-1"
        )
    )
    print(f"索引 '{index_name}' 创建成功。")
    
    # 连接到索引
    index = pc.Index(index_name)
    print(f"已成功连接到索引 '{index_name}'")
    
    return index

# 3. 加载MNIST数据并插入到Pinecone
def load_and_insert_data(index):
    """加载MNIST数据并插入到Pinecone索引"""
    
    # 加载MNIST数据集
    digits = load_digits(n_class=10)  # 加载0-9所有数字
    
    # 获取数据集中的特征数据和标签
    X = digits.data  # 64维向量（8×8像素展平）
    y = digits.target  # 数字标签（0-9）
    
    print(f"数据集大小: {len(X)} 个样本")
    print(f"特征维度: {X.shape[1]} 维")
    print(f"标签类别: {set(y)}")
    
    # 初始化向量列表
    vectors = []
    
    # 转换数据格式
    for i in range(len(X)):
        vector_id = str(i)  # 使用样本索引作为向量ID
        vector_values = X[i].tolist()  # 转换为Python列表
        metadata = {'label': int(y[i])}  # 存储标签信息
        
        vectors.append((vector_id, vector_values, metadata))
    
    # 分批插入数据（Pinecone每批次最多1000个向量）
    batch_size = 100
    total_vectors = len(vectors)
    
    print(f"开始插入数据，共 {total_vectors} 个向量，批大小: {batch_size}")
    
    for i in range(0, total_vectors, batch_size):
        batch = vectors[i:i + batch_size]
        index.upsert(vectors=batch)
        print(f"已插入 {min(i + batch_size, total_vectors)}/{total_vectors} 个向量")
    
    print("所有数据插入完成")
    
    return X, y, digits

# 4. 相似性搜索
def similarity_search(index, query_vector, top_k=5):
    """在Pinecone索引中执行相似性搜索"""
    
    results = index.query(
        vector=query_vector,
        top_k=top_k,
        include_metadata=True
    )
    
    return results

# 5. 创建测试数字图像
def create_test_digit(digit_value=3):
    """创建手写数字测试图像"""
    
    # 根据数字值创建不同的测试图像
    digit_patterns = {
        0: np.array([
            [0, 255, 255, 255, 255, 255, 255, 0],
            [255, 255, 0, 0, 0, 0, 255, 255],
            [255, 255, 0, 0, 0, 255, 255, 255],
            [255, 255, 0, 0, 255, 255, 255, 255],
            [255, 255, 0, 255, 255, 0, 255, 255],
            [255, 255, 255, 255, 0, 0, 255, 255],
            [255, 255, 255, 0, 0, 0, 255, 255],
            [0, 255, 255, 255, 255, 255, 255, 0]
        ]),
        1: np.array([
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 255, 0, 0, 0],
            [0, 255, 255, 255, 255, 0, 0, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 255, 255, 255, 255, 255, 255, 0]
        ]),
        2: np.array([
            [0, 255, 255, 255, 255, 255, 255, 0],
            [255, 255, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 0, 0, 255, 255, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 0, 0, 0, 0],
            [0, 255, 255, 0, 0, 0, 0, 0],
            [255, 255, 255, 255, 255, 255, 255, 255]
        ]),
        3: np.array([
            [0, 255, 255, 255, 255, 255, 255, 0],
            [255, 255, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 255, 255, 255, 255, 0],
            [0, 0, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 0, 0, 0, 255, 255],
            [255, 255, 0, 0, 0, 0, 255, 255],
            [0, 255, 255, 255, 255, 255, 255, 0]
        ]),
        7: np.array([
            [255, 255, 255, 255, 255, 255, 255, 255],
            [0, 0, 0, 0, 0, 0, 255, 255],
            [0, 0, 0, 0, 0, 255, 255, 0],
            [0, 0, 0, 0, 255, 255, 0, 0],
            [0, 0, 0, 255, 255, 0, 0, 0],
            [0, 0, 255, 255, 0, 0, 0, 0],
            [0, 255, 255, 0, 0, 0, 0, 0],
            [255, 255, 0, 0, 0, 0, 0, 0]
        ])
    }
    
    # 如果指定的数字没有预定义模式，使用默认模式
    if digit_value not in digit_patterns:
        digit_value = 3
    
    digit_image = digit_patterns[digit_value]
    
    # 将图像像素值从0~255的范围缩放到0~16的范围
    # 匹配MNIST数据集中使用的像素值范围
    digit_flatten = (digit_image / 255.0) * 16
    
    # 将二维图像数组展开成一维列表
    query_data = digit_flatten.ravel().tolist()
    
    return digit_image, query_data, digit_value

# 6. 可视化结果
def visualize_results(query_image, results, true_label, predicted_label):
    """可视化查询图像和搜索结果"""
    
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # 显示查询图像
    axes[0].imshow(query_image, cmap='gray')
    axes[0].set_title(f'查询数字 (真实: {true_label}, 预测: {predicted_label})', size=12)
    axes[0].axis('off')
    
    # 显示搜索结果信息
    axes[1].text(0.1, 0.9, "最近邻搜索结果:", fontsize=12, fontweight='bold')
    
    labels = []
    for i, match in enumerate(results['matches']):
        label = match['metadata']['label']
        distance = match['score']
        labels.append(label)
        axes[1].text(0.1, 0.8 - i*0.1, f"结果 {i+1}: 数字 {label}, 距离: {distance:.4f}", fontsize=10)
    
    # 显示投票结果
    vote_counts = Counter(labels)
    most_common = vote_counts.most_common(1)[0]
    axes[1].text(0.1, 0.3, f"投票结果: 数字 {most_common[0]} (出现{most_common[1]}次)", 
                fontsize=11, fontweight='bold', color='red')
    
    axes[1].set_xlim(0, 1)
    axes[1].set_ylim(0, 1)
    axes[1].axis('off')
    
    plt.tight_layout()
    plt.show()

# 主函数
def main():
    # 你的Pinecone API密钥
    API_KEY = "pcsk_5XCTf3_7o7fbVgtePSb8P89EiDJAUNpX2Dc1Dan2GnuML5WQoWTcznFvE5dVhFTmd6A6MC"
    INDEX_NAME = "mnist-index"
    
    print("=== Pinecone MNIST 数字识别项目 ===")
    
    # 1. 初始化Pinecone
    pc = initialize_pinecone(API_KEY)
    if pc is None:
        return
    
    # 2. 创建并连接索引
    index = create_and_connect_index(pc, INDEX_NAME)
    
    # 3. 加载并插入数据
    X, y, digits = load_and_insert_data(index)
    
    # 4. 测试不同的数字
    test_digits = [3, 1, 2, 7, 0]
    
    for digit in test_digits:
        print(f"\n=== 测试数字 {digit} ===")
        
        # 创建测试图像
        query_image, query_vector, true_label = create_test_digit(digit)
        
        # 执行相似性搜索
        results = similarity_search(index, query_vector, top_k=11)
        
        # 提取标签并进行投票
        labels = [match['metadata']['label'] for match in results['matches']]
        final_prediction = Counter(labels).most_common(1)[0][0]
        
        print(f"真实数字: {true_label}")
        print(f"预测数字: {final_prediction}")
        print(f"准确率: {true_label == final_prediction}")
        
        # 显示详细搜索结果
        print("\n最近邻搜索结果:")
        for i, match in enumerate(results['matches']):
            label = match['metadata']['label']
            distance = match['score']
            print(f"  结果 {i+1}: 数字 {label}, 距离: {distance:.4f}")
        
        # 可视化结果
        visualize_results(query_image, results, true_label, final_prediction)

if __name__ == "__main__":
    main()